//===============================================================================
// Microsoft Architecture Strategy Team
// LitwareHR - SaaS Sample Application
//===============================================================================
// Copyright  Microsoft Corporation.  All rights reserved.
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY
// OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT
// LIMITED TO THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
// FITNESS FOR A PARTICULAR PURPOSE.
//===============================================================================
// The example companies, organizations, products, domain names,
// e-mail addresses, logos, people, places, and events depicted
// herein are fictitious.  No association with any real company,
// organization, product, domain name, email address, logo, person,
// places, or events is intended or should be inferred.
//===============================================================================

using System;
using System.Collections.ObjectModel;
using System.IdentityModel.Tokens;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Security.Tokens;
using System.Xml;
using System.Globalization;
using System.Collections.Generic;

namespace Shp.Security.BrokeredReceiver
{
    /// <summary>
    /// Abstract base class for STS implementations
    /// </summary>
	public abstract class SecurityTokenService : ISecurityTokenService
	{
		string stsName; // The name of the STS. Used to populate saml:Assertion/@Issuer
        SecurityToken issuerToken; // The SecurityToken used to sign issued tokens
        Dictionary<Uri, X509SecurityToken> proofKeyEncryptionTokens; // The collection of SecurityTokens used to encrypt the proof key in the issued token.

        /// <summary>
        /// constructor
        /// </summary>
        /// <param name="stsName">The name of the STS. Used to populate saml:Assertion/@Issuer</param>
        /// <param name="issuerToken">The issuer token.</param>
        /// <param name="encryptionTokens">The encryption tokens.</param>
        protected SecurityTokenService(string stsName, X509SecurityToken issuerToken, 
            Dictionary<Uri, X509SecurityToken> encryptionTokens)
        {
            this.stsName = stsName;
            this.issuerToken = issuerToken;
            this.proofKeyEncryptionTokens = encryptionTokens;
        }

        /// <summary>
        /// The name of the STS.
        /// </summary>
        protected string SecurityTokenServiceName
        {
            get { return this.stsName; }
        }


        /// <summary>
        /// The SecurityToken used to sign tokens the STS issues.
        /// </summary>
        protected SecurityToken IssuerToken
        {
            get { return this.issuerToken; }
        }

        /// <summary>
        /// The collection of SecurityTokens used to encrypt the proof key in the issued token.
        /// </summary>
        protected Dictionary<Uri, X509SecurityToken> ProofKeyEncryptionTokens
        {
            get { return this.proofKeyEncryptionTokens; }
        }

		#region Abstract methods
        

        /// <summary>
        /// abstract method for setting up claims in the SAML Token issued by the STS
        /// Should be overriden by STS implementations deriving from this base class
        /// to set up appropriate claims
        /// </summary>
        protected abstract Collection<SamlAttribute> GetIssuedClaims(RequestSecurityToken requestSecurityToken);

        #endregion

        /// <summary>
        /// Virtual method for ProcessRequestSecurityToken
        /// Should be overriden by STS implementations deriving from this base class
        /// </summary>
        public virtual Message IssueSecurityToken(Message message)
        {
            // Check for appropriate action header
            EnsureRequestSecurityTokenAction(message);

            // Extract the MessageID from the request message
            UniqueId requestMessageID = message.Headers.MessageId;

            if (requestMessageID == null)
            {
                FederationUtilities.ThrowTrustFault(Properties.Resources.MissingMessageId);
            }

            RequestSecurityToken rst = RequestSecurityToken.CreateFrom(message.GetReaderAtBodyContents());
            
            // Set up the claims we are going to issue
            Collection<SamlAttribute> samlAttributes = GetIssuedClaims(rst);

            int keySize = (rst.KeySize != 0) ? rst.KeySize : 192;

            // Create proof token
            // Get requestor entropy, if any
            byte[] senderEntropy = null;
            SecurityToken entropyToken = rst.RequestorEntropy;
            if (entropyToken != null)
            {
                senderEntropy = ((BinarySecretSecurityToken)entropyToken).GetKeyBytes();
            }

            byte[] key = null;
            byte[] stsEntropy = null;

            // If sender provided entropy, then we're using combined entropy
            if (senderEntropy != null)
            {
                stsEntropy = new byte[keySize / 8];
                RNGCryptoServiceProvider random = new RNGCryptoServiceProvider();
                random.GetNonZeroBytes(stsEntropy);
                key = RequestSecurityTokenResponse.ComputeCombinedKey(senderEntropy, stsEntropy, keySize);
            }
            else // Issuer entropy only...
            {
                key = new byte[keySize / 8];
                RNGCryptoServiceProvider random = new RNGCryptoServiceProvider();
                random.GetNonZeroBytes(key);
            }

            // Create a BinarySecretSecurityToken to be the proof token, based on the key material
            // in key. key will be the combined key in the combined entropy case, or the issuer entropy
            // otherwise
            BinarySecretSecurityToken proofToken = new BinarySecretSecurityToken(key);

            // Create a SAML token, valid for around 10 hours
            DateTime notBefore = DateTime.UtcNow - TimeSpan.FromMinutes(5);
            DateTime notOnOrAfter = DateTime.UtcNow + TimeSpan.FromHours(10);
            SamlSecurityToken samlToken = SamlTokenCreator.CreateSamlToken(this.stsName,
                                                                           proofToken,
                                                                           this.IssuerToken,
                                                                           SelectProofKeyEncryptionToken(rst.AppliesTo.Uri),
                                                                           new SamlConditions(notBefore, notOnOrAfter),
                                                                           samlAttributes);

            RequestSecurityTokenBase rstr = GetRequestSecurityTokenResponse(
                rst, keySize, proofToken, samlToken, senderEntropy, stsEntropy);   
        
            Message rstrMessage = Message.CreateMessage(
                message.Version, Constants.TrustFeb2005.RequestSecurityTokenIssuanceResponse, rstr);
            rstrMessage.Headers.RelatesTo = requestMessageID;

            return rstrMessage;
        }

        #region Helper Methods

        protected SecurityToken SelectProofKeyEncryptionToken(Uri target)
        {
            SecurityToken token = this.ProofKeyEncryptionTokens[target];
            if (token == null)
            {
                FederationUtilities.ThrowTrustFault(Properties.Resources.EncryptionTokenNotFound, target.AbsoluteUri);
            }
            return token;
        }

        /// <summary>
        /// Validate action header and discard messages with inappropriate headers
        /// </summary>
        protected static void EnsureRequestSecurityTokenAction(Message message)
        {
            if (message == null)
            {
                FederationUtilities.ThrowTrustFault(Properties.Resources.ArgumentNullException, "message");
            }

            if (message.Headers.Action != Constants.TrustFeb2005.RequestSecurityTokenIssuance)
            {
                FederationUtilities.ThrowTrustFault(Properties.Resources.UnsupportedActionException, message.Headers.Action);
            }
        }

        /// <summary>
        /// Helper Method to Create Proof Token. Creates BinarySecretSecuryToken 
        /// with the requested number of bits of random key material
        /// </summary>
        /// <param name="keySize">keySize</param>
        /// <returns>Proof Token</returns>
        protected static BinarySecretSecurityToken CreateProofToken(int keySize)
        {
            byte[] key = new byte[keySize / 8];
            RNGCryptoServiceProvider random = new RNGCryptoServiceProvider();
            random.GetNonZeroBytes(key);
            return new BinarySecretSecurityToken(key);
        }

        /// <summary>
        /// Helper Method to set up the RSTR
        /// </summary>
        /// <param name="rst">RequestSecurityToken</param>
        /// <param name="keySize">keySize</param>
        /// <param name="proofToken">proofToken</param>
        /// <param name="samlToken">The SAML Token to be issued</param>
        /// <returns>RequestSecurityTokenResponse</returns>
        protected static RequestSecurityTokenBase GetRequestSecurityTokenResponse(RequestSecurityTokenBase requestSecurityToken,
                                                                                      int keySize,
                                                                                      SecurityToken proofToken,
                                                                                      SecurityToken samlToken,
                                                                                      byte[] senderEntropy,
                                                                                      byte[] stsEntropy)
        {
            // Create an uninitialized RequestSecurityTokenResponse object and set the various properties
            RequestSecurityTokenResponse rstr = new RequestSecurityTokenResponse();
            rstr.TokenType = Constants.SecurityTokenService.SamlTokenTypeUri;
            rstr.RequestedSecurityToken = samlToken;
            rstr.RequestedUnattachedReference = samlToken.CreateKeyIdentifierClause<SamlAssertionKeyIdentifierClause>();
            rstr.RequestedAttachedReference = samlToken.CreateKeyIdentifierClause<SamlAssertionKeyIdentifierClause>();
            rstr.Context = requestSecurityToken.Context;
            rstr.KeySize = keySize;

            // If sender provided entropy then we'll be using combined entropy so set the IssuerEntropy
            if (senderEntropy != null)
            {
                rstr.IssuerEntropy = new BinarySecretSecurityToken(stsEntropy);
                rstr.ComputeKey = true;
            }
            else // Issuer entropy only...
            {
                rstr.RequestedProofToken = proofToken;
            }

            return rstr;
        }

        #endregion
    }
}
